Bayesian Data Selection

Insights into complex, high-dimensional data can be obtained by discovering features of the data that match or do not match a model of interest. To formalize this task, we introduce the "data selection" problem: finding a lower-dimensional statistic—such as a subset of variables—that is well fit by a given parametric model of interest. A fully Bayesian approach to data selection would be to parametrically model the value of the statistic, nonparametrically model the remaining "background" components of the data, and perform standard Bayesian model selection for the choice of statistic. However, fitting a nonparametric model to high-dimensional data tends to be highly inefficient, statistically and computationally. We propose a novel score for performing data selection, the "Stein volume criterion (SVC)", that does not require fitting a nonparametric model. The SVC takes the form of a generalized marginal likelihood with a kernelized Stein discrepancy in place of the Kullback–Leibler divergence. We prove that the SVC is consistent for data selection, and establish consistency and asymptotic normality of the corresponding generalized posterior on parameters. We apply the SVC to the analysis of single-cell RNA sequencing data sets using probabilistic principal components analysis and a spin glass model of gene regulation.


Introduction
Scientists often seek to understand complex phenomena by developing working models for various special cases and subsets. Thus, when faced with a large complex dataset, a natural question to ask is where and when a given working model applies. We formalize this question statistically by saying that given a high-dimensional dataset, we want to identify a lowerdimensional statistic-such as a subset of variables-that follows a parametric model of interest (the working model). We refer to this problem as "data selection", in counterpoint to model selection, since it requires selecting the aspect of the data to which a given model applies.
For example, early studies of single-cell RNA expression showed that the expression of individual genes was often bistable, which suggests that the system of cellular gene expression might be described with the theory of interacting bistable systems, or spin glasses, with each gene a separate spin and each cell a separate observation. While it seems implausible that such a model would hold in full generality, it is quite possible that there are subsets of genes for which the spin glass model is a reasonable approximation to reality. Finding such subsets of genes is a data selection problem. In general, a good data selection method would enable one to (a) discover interesting phenomena in complex datasets, (b) identify precisely where naive application of the working model to the full dataset goes wrong, and (c) evaluate the robustness of inferences made with the working model.
Perhaps the most natural Bayesian approach to data selection is to employ a semiparametric joint model, using the parametric model of interest for the low-dimensional statistic (the "foreground") and using a flexible nonparametric model to explain all other aspects of the data (the "background"). Then, to infer where the foreground model applies, one would perform standard Bayesian model selection across different choices of the foreground statistic. However, this is computationally challenging due to the need to integrate over the nonparametric model for each choice of foreground statistic, making this approach quite difficult in practice. A natural frequentist approach to data selection would be to perform a goodness-of-fit test for each choice of foreground statistic. However, this still requires specifying an alternative hypothesis, even if the alternative is nonparametric, and ensuring comparability between alternatives used for different choices of foreground statistics is nontrivial. Moreover, developing goodness-of-fit tests for composite hypotheses or hierarchical models is often difficult in practice.
In this article, we propose a new score-for both data selection and model selectionthat is similar to the marginal likelihood of a semi-parametric model but does not require one to specify a background model, let alone integrate over it. The basic idea is to employ a generalized marginal likelihood where we replace the foreground model likelihood by an exponentiated divergence with nice properties, and replace the background model's marginal likelihood with a simple volume correction factor. For the choice of divergence, we use a kernelized Stein discrepancy (KSD) since it enables us to provide statistical guarantees and is easy to estimate compared to other divergences -for instance, the Kullback-Leibler divergence involves a problematic entropy term that cannot simply be dropped. The background model volume correction arises roughly as follows: if the background model is well-specified, then asymptotically, its divergence from the empirical distribution converges to zero and all that remains of the background model's contribution is the volume of its effective parameter space. Consequently, it is not necessary to specify the background model, only its effective dimension. To facilitate computation further, we develop a Laplace approximation for the foreground model's contribution to our proposed score.
This article makes a number of novel contributions. We introduce the data selection problem in broad generality, and provide a thorough asymptotic analysis. We propose a novel model/data selection score, which we refer to as the Stein volume criterion, that takes the form of a generalized marginal likelihood using a KSD. We provide new theoretical results for this generalized marginal likelihood and its associated posterior, complementing and building upon recent work on the frequentist properties of minimum KSD estimators . Finally, we provide first-of-a-kind empirical data selection analyses with two models that are frequently used in single-cell RNA sequencing analysis.
The article is organized as follows. In Section 2, we introduce the data selection problem and our proposed method. In Section 3 we study the asymptotic properties of Bayesian data selection methods and compare to model selection. Section 4 provides a review of related work and Section 5 illustrates the method on a toy example. In Section 6, we prove (a) consistency results for both data selection and model selection, (b) a Laplace approximation for the proposed score, and (c) a Bernstein-von Mises theorem for the corresponding generalized posterior. In Section 7, we apply our method to probabilistic principal components analysis (pPCA), assess its performance in simulations, and demonstrate it on single-cell RNA sequencing (scRNAseq) data. In Section 8, we apply our method to a spin glass model of gene expression, also demonstrated on an scRNAseq dataset. Section 9 concludes with a brief discussion.

Method
Suppose the data X (1) , . . . , X (N ) ∈ X are independent and identically distributed (i.i.d.), where X ⊆ R d . Suppose the true data-generating distribution P 0 has density p 0 (x) with respect to Lebesgue measure, and let {q(x|θ) : θ ∈ Θ} be a parametric model of interest, where Θ ⊆ R m . We are interested in evaluating this model when applied to a projection of the data onto a subspace, X F ⊆ X (the "foreground" space). Specifically, let X F := V X be a linear projection of X ∈ X onto X F , where V is a matrix with orthonormal columns. Let q(x F |θ) denote the distribution of X F when X ∼ q(x|θ), and likewise, let p 0 (x F ) be the distribution of X F when X ∼ p 0 (x). Even when the complete model q(x|θ) is misspecified with respect to p 0 (x), it may be that q(x F |θ) is well-specified with respect to p 0 (x F ); see Figure 1 for a toy example. In such cases, the parametric model is only partially misspecified -specifically, it is misspecified on the "background" space X B , defined as the orthogonal complement of X F . Our goal is to find subspaces X F for which q(x F |θ) is correctly specified.
A natural Bayesian solution would be to replace the background component of the assumed model, q(x B |x F , θ), with a more flexible componentq(x B |x F , φ B ) that is guaranteed to be well-specified with respect to p 0 (x B |x F ), such as a nonparametric model. The resulting joint model, which we refer to as the "augmented model", is then θ ∼ π(θ), X  , φ B )π(θ)π B (φ B )dθdφ B . However, in general, it is difficult to find a background model that (a) is guaranteed to be well-specified with respect to p 0 (x B |x F ) and (b) can be integrated over in a computationally tractable way to obtain the posterior on the choice of F. Our proposed method, which we introduce next, sidesteps these difficulties while still exhibiting similar guarantees.

Proposed score for data selection and model selection
In this section, we propose a model/data selection score that is simpler to compute than the marginal likelihood of the augmented model and has similar theoretical guarantees. This score takes the form of a generalized marginal likelihood with a normalized kernelized Stein discrepancy (nksd) estimate taking the place of the log likelihood. Specifically, our proposed model/data selection score, termed the "Stein volume criterion" (SVC), is where the "temperature" T > 0 is a hyperparameter and m B is the effective dimension of the background model parameter space. nksd(· ·) is an empirical estimate of the nksd; see Equations 4 and 5. The integral in Equation 2 can be approximated using techniques discussed in Section 2.3. The hyperparameter T can be calibrated by comparing the coverage of the standard Bayesian posterior to the coverage of the nksd generalized posterior (Section S1.1). The (2π/N ) m B /2 factor penalizes higher-complexity background models. In general, we allow m B to grow with N , particularly when the background model is nonparametric. Crucially, the likelihood of the background model does not appear in our proposed score, sidestepping the need to fit or even specify the background model -indeed, the only place that the background model enters into the SVC is through m B . Thus, rather than specify a background model and then derive m B , one can simply specify an appropriate value of m B . Reasonable choices of m B can be derived by considering the asymptotic behavior of a Pitman-Yor process mixture model, a common nonparametric model that is a natural choice for a background model. A Pitman-Yor process mixture model with discount parameter α ∈ (0, 1), concentration parameter θ > −α, and D-dimensional component parameters will asymptotically have expected effective dimension under the prior, where a N ∼ b N means that a N /b N → 1 as N → ∞ and Γ(·) is the gamma function (Pitman, 2002, §3.3). As a default, we recommend setting m B = c B r B √ N , where r B is the dimension of X B and c B is a constant chosen to match Equation 3 with α = 1/2. The √ N scaling is particularly nice in terms of asymptotic guarantees; see Section 3.2. The SVC uses a novel, normalized version of the ksd between densities p(x) and q(x): where k(x, y) ∈ R is an integrally strictly positive definite kernel, s q (x) := ∇ x log q(x), and s p (x) := ∇ x log p(x); see Section 6.1 for details. The numerator corresponds to the standard ksd . The denominator, which is strictly positive and independent of q(x), is a normalization factor that we have introduced to make the divergence comparable across spaces of different dimension. See Section S1.2 for kernel recommendations. Extending the technique of , we propose to estimate the normalized KSD using U-statistics: where X (i) ∼ p(x) i.i.d., the sums are over all i, j ∈ {1, . . . , N } such that i = j, and Importantly, Equation 5 does not require knowledge of s p (x), which is unknown in practice.

Comparison with the standard marginal likelihood
It is instructive to compare our proposed model/data selection score, the Stein volume criterion, to the standard marginal likelihoodq(X (1:N ) |F). In particular, we show that the SVC approximates a generalized version of the marginal likelihood. To see this, first define H := − p 0 (x) log p 0 (x)dx, the entropy of the complete data distribution, and note that if were H somehow known, then the Kullback-Leibler (kl) divergence between the augmented model and the data distribution could be approximated as Since multiplying the marginal likelihoods by a fixed constant does not affect the Bayes factors, the following expression could be used instead of the marginal likelihoodq(X (1:N ) |F) to decide among foreground subspaces: Now, consider a generalized marginal likelihood where the nksd replaces the kl: We refer toK as the "nksd marginal likelihood" of the augmented model. Intuitively, we expect it to behave similarly to the standard marginal likelihood, except that it quantifies the divergence between the model and data distributions using the nksd instead of the kl. However, a key advantage of the nksd marginal likelihood is that it admits a simple approximation via the SVC when the background model is well-specified, unlike the standard marginal likelihood. For instance, if the foreground and background are independent, that is, , then the theory in Section 6 can be extended to the full augmented model to show that where K is the SVC (Equation 2). Thus, the SVC approximates the nksd marginal likelihood of the augmented model, suggesting that the SVC may be a convenient alternative to the standard marginal likelihood. Formally, Section 3 shows that the SVC exhibits consistency properties similar to the standard marginal likelihood, even when p 0 (x) = p 0 (x F )p 0 (x B ).

Computation
Next, we discuss methods for computing the SVC including exact solutions, Laplace/BIC approximation, variational approximation, and comparing many possible choices of F.

Exact solution for exponential families
When the foreground model is an exponential family, the SVC can be computed analytically. Specifically, in Section S1.3, we show if q( where A, B, and C depend on the data X (1:N ) but not on θ. Therefore, we can place a multivariate Gaussian prior on θ and compute the SVC in closed form; see Section S1.3.

Laplace and BIC approximations
The Laplace approximation is a widely-used technique for computing marginal likelihoods. In Theorem 6.9, we establish regularity conditions under which a Laplace approximation to the SVC is justified by being asymptotically correct. The resulting approximation is where θ N := argmin θ nksd(p 0 (x F ) q(x F |θ)) is the point at which the estimated nksd is minimized, the "minimum Stein discrepancy estimator" as defined by . We can also make a rougher approximation, analogous to the Bayesian information criterion (BIC), which does not require one to compute second derivatives of nksd: This approximation is easy to compute, given a minimum Stein discrepancy estimator θ N . Like the SVC, it satisfies all of our consistency desiderata (Section S2). However, we expect it to perform worse than the SVC when there is not yet enough data for the nksd posterior to be highly concentrated, that is, when a range of θ values can plausibly explain the data.

Comparing many foregrounds using approximate optima
Often, we would like to evaluate many possible subspaces X F when performing data selection. Even when using the Laplace or BIC approximation to the SVC, this can get computationally prohibitive since we need to re-optimize to find θ N for every F under consideration. Here, we propose a way to reduce this cost by making a fast linear approximation. Define j (θ) := nksd(p 0 (x F j ) q(x F j |θ)) for j ∈ {1, 2}. For w ∈ [0, 1], we can linearly interpolate Now, θ N (0) and θ N (1) are the minimum Stein discrepancy estimators for F 1 and F 2 , respectively. Given θ N (0), we can approximate θ N (1) by applying the implicit function theorem and a first-order Taylor expansion (Section S1.4): Note that the derivatives of j are often easy to compute with automatic differentiation (Baydin et al., 2018). Note also that when we are comparing one foreground subspace, such as X F 1 = X , to many other foreground subspaces X F 2 , the inverse Hessian ∇ 2 θ 1 (θ N (0)) −1 only needs to be computed once. Thus, Equation 13 provides a fast method for computing Laplace or BIC approximations to the SVC for a large number of candidate foregrounds F.

Variational approximation
Variational inference is a method for approximating both the posterior distribution and the marginal likelihood of a probabilistic model. Since the SVC takes the form of a generalized marginal likelihood, we can derive a variational approximation to the SVC. Let r ζ (θ) be an approximating distribution parameterized by ζ. By Jensen's inequality, we have Maximizing this lower bound with respect to the variational parameters ζ provides an approximation to the SVC, or more precisely, to log K − (m B /2) log(2π/N ). Note that this variational approximation falls within the framework of generalized variational inference proposed by Knoblauch et al. (2019).

Data selection and model selection consistency
This section presents our consistency results when comparing two different foreground subspaces (data selection) or two different foreground models (model selection). The theory supporting these results is in Sections 6 and S2. We consider four distinct properties that a procedure would ideally exhibit: data selection consistency, nested data selection consistency, model selection consistency, and nested model selection consistency; see Section 6.4 for precise definitions. We consider six possible model/data selection scores, and we establish which scores satisfy which properties; see Table 1. The SVC and the full marginal likelihood are the only two of the six scores that satisfy all four consistency properties.
The intuition behind Bayesian model selection is often explained in terms of Occam's razor: a theory should be as simple as possible but no simpler. Data selection and nested data selection encapsulate a complementary intuition: a theory should explain as much of the data as possible but no more. In other words, when choosing between foreground spaces, a consistent data selection score will asymptotically prefer the highest-dimensional space on which the model is correctly specified. As in standard model selection, a practical concern in data selection is robustness. For instance, if the foreground model is even slightly misspecified on X F 2 , then the empty foreground X F 1 = ∅ will be asymptotically preferred over X F 2 . Since the SVC takes the form of a generalized marginal likelihood, techniques for improving robustness with the standard marginal likelihood-such as coarsened posteriors, power posteriors, and BayesBag-could potentially be extended to address this issue (Miller and Dunson, 2019;Huggins and Miller, 2021). We leave exploration of such approaches to future work.

Data selection consistency
First, consider comparisons between different choices of foreground, F 1 and F 2 . When the model is correctly specified over F 1 but not F 2 , we refer to asymptotic concentration on F 1 as "data selection consistency" (and vice versa if F 2 is correct but not F 1 ). For the standard marginal likelihood of the augmented model, we have (see Section S2.2) where θ kl j, * := argmin kl(p 0 (x F j ) q(x F j |θ)) for j ∈ {1, 2}, that is, θ kl j, * is the parameter value that minimizes the kl divergence between the projected data distribution p 0 (x F j ) and the projected model q(x F j |θ). Thus,q(X (1:N ) |F j ) asymptotically concentrates on the F j on which the projected model can most closely match the data distribution in terms of kl.
In Theorem 6.17, we show that under mild regularity conditions, the Stein volume criterion behaves precisely the same way but with the nksd in place of the kl: where θ nksd j, * := argmin nksd(p 0 (x F j ) q(x F j |θ)) for j ∈ {1, 2}. Therefore,q(X (1:N ) |F) and K both yield data selection consistency. It is important here that the SVC uses a true divergence, rather than a divergence up to a data-dependent constant. If we instead used which employs the foreground marginal likelihood q(X |θ)π(θ)dθ and a background volume correction, we would get qualitatively different behavior (Section S2.2): In short, the naive score K (a) is a bad choice: it decides between data subspaces based not just on how well the parametric foreground model performs, but also on the entropy of the data distribution in each space. As a result, K (a) does not exhibit data selection consistency.

Nested data selection consistency
When X F 2 ⊂ X F 1 , we refer to the problem of deciding between subspaces F 1 and F 2 as nested data selection, in counterpoint to nested model selection, where one model is a subset of another (Vuong, 1989). If the model q(x|θ) is well-specified over X F 1 , then it is guaranteed to be well-specified over any lower-dimensional sub-subspace X F 2 ⊂ X F 1 ; in this case, we refer to asymptotic concentration on F 1 as "nested data selection consistency". In this situation, kl(p 0 (x F j ) q(x F j |θ kl j, * )) and nksd(p 0 (x F j ), q(x F j |θ nksd j, * )) are both zero for j ∈ {1, 2}, making it necessary to look at higher-order terms in Equations 15 and 16. In Section S2.3, we show that if X F 2 ⊂ X F 1 , q(x|θ) is well-specified over X F 1 , the background models are well-specified, and their dimensions m B 1 and m B 2 are constant with respect to N , then where m F j is the effective dimension of the parameter space of q(x F j |θ). In Theorem 6.17, we show that under mild regularity conditions, the SVC behaves the same way: Thus, so long as m F 2 + m B 2 > m F 1 + m B 1 whenever X F 2 ⊂ X F 1 , the marginal likelihood and the SVC asymptotically concentrate on the larger foreground F 1 ; hence, they both exhibit nested data selection consistency. This is a natural assumption since the background model is generally more flexible-on a per dimension basis-than the foreground model. The volume correction (2π/N ) m B /2 in the definition of the SVC is important for nested data selection consistency (Equation 20). An alternative score without that correction, exhibits data selection consistency (Equation 16 holds for K (b) ), but not nested data selection consistency; see Sections S2.2 and S2.3. More subtly, the asymptotics of the SVC in the case of nested data selection also depend on the variance of U-statistics. To illustrate, consider a score that is similar to the SVC but uses kl instead of nksd: F |θ) − H F and H F is required to be known. The score K (c) exhibits data selection consistency, but not nested data selection consistency. The reason is that the error in estimating the kl is of order 1/ √ N by the central limit theorem, and this source of error dominates the log N term contributed by the volume correction; see Section S2.3. Meanwhile, the error in estimating the nksd is of order 1/N when the model is well-specified, due to the rapid convergence rate of the U-statistic estimator. Thus, in the SVC, this source of error is dominated by the volume correction; see Theorem 6.12.
The nested data selection results we have described so far assume m B does not depend on N , or at least m B 2 − m B 1 does not depend on N (Theorem 6.17). However, in Section 2.1, we suggest setting m B = c B r B √ N where c B is a constant and r B is the dimension of X B . With this choice, the asymptotics of the SVC for nested data selection become (Theorem 6.17) Since r B 1 < r B 2 when X F 2 ⊂ X F 1 , the SVC concentrates on the larger foreground F 1 , yielding nested data selection consistency. Going beyond the well-specified case, Theorem 6.17 shows that Equation 23 holds when nksd(p 0 (x F 1 ) q(x F 1 |θ nksd 1, * )) = nksd(p 0 (x F 2 ) q(x F 2 |θ nksd 2, * )) = 0, that is, when the models are misspecified by the same amount as measured by the nksd. Equation 23 holds regardless of whether m F 1 is equal to m F 2 .

Model selection and nested model selection consistency
Consider comparing different foreground models q 1 (x F |θ 1 ) and q 2 (x F |θ 2 ) over the same subspace X F , while using the same background model. We say that a score exhibits "model selection consistency" if it concentrates on the correct model, when one of the models is correctly specified and the other is not. When the two models are nested and both are correct, a score exhibits "nested model selection consistency" if it concentrates on the simpler model.
Like the standard marginal likelihood, the SVC exhibits both types of model selection consistency. The standard marginal likelihood satisfies (Section S2.4) where θ kl j, * := argmin kl(p 0 (x F ) q j (x F |θ j )) for j ∈ {1, 2}. Analogously, by Theorem 6.17, where θ nksd j, * := argmin nksd(p 0 (x F ) q j (x F |θ j )) for j ∈ {1, 2}. Thus, for both scores, concentration occurs on the model that comes closer to the data distribution in terms of the corresponding divergence (kl or nksd).
For nested model selection, suppose both foreground models are well-specified and m B 1 = m B 2 . Letting m F ,j be the parameter dimension of q j (x F |θ j ), we have (Section S2.5) In Theorem 6.17, we show that the SVC behaves identically: Here, a key role is played by the volume of the foreground parameter space, which quantifies the foreground model complexity. The SVC accounts for this by integrating over foreground parameter space. Meanwhile, a naive alternative that ignores the foreground volume, exhibits model selection consistency (Equation 25 holds for K (d) ) but not nested model selection consistency (Section S2.5). The Laplace and BIC approximations to the SVC (Equations 10 and 11) explicitly correct for the foreground parameter volume without integrating.

Related work
Projection pursuit methods are closely related to data selection in that they attempt to identify "interesting" subspaces of the data. However, projection pursuit uses certain prespecified objective functions to optimize over projections, whereas our method allows one to specify a model of interest (Huber, 1985). Another related line of research is on Bayesian goodness-of-fit (GOF) tests, which compute the posterior probability that the data comes from a given parametric model versus a flexible alternative such as a nonparametric model. Our setup differs in that it aims to compare among different semiparametric models. Nonetheless, in an effort to address the GOF problem, a number of authors have developed nonparametric models with tractable marginals (Verdinelli and Wasserman, 1998;, and using these models as the background component in an augmented model could in theory solve data selection problems. In practice, however, such models can only be applied to one-dimensional or few-dimensional data spaces. In Section 7, we show that naively extending the method of  to the multi-dimensional setting has fundamental limitations. There is a sizeable frequentist literature on GOF testing using discrepancies (Gretton et al., 2012;Barron, 1989;Györfi and Van Der Meulen, 1991). Our proposed method builds directly on the KSD-based GOF test proposed by  and Chwialkowski et al. (2016). However, using these methods to draw comparisons between different foreground subspaces is non-trivial, since the set of alternative models considered by the GOF test, though nonparametric, will be different over data spaces with different dimensionality. Moreover, the Bayesian aspect of the SVC makes it more straightforward to integrate prior information and employ hierarchical models.
In composite likelihood methods, instead of the standard likelihood, one uses the product of the conditional likelihoods of selected statistics (Lindsay, 1988;Varin et al., 2011). Composite likelihoods have seen widespread use, often for robustness or computational purposes. However, in composite likelihood methods, the choice of statistics is fixed before performing inference. In contrast, in data selection the choice of statistics is a central quantity to be inferred.
Relatedly, our work connects with the literature on robust Bayesian methods. Doksum and Lo (1990) propose conditioning on the value of an insufficient statistic, rather than the complete dataset, when performing inference; also see Lewis et al. (2021). However, making an appropriate choice of statistic requires one to know which aspects of the model are correct; in contrast, our procedure infers the choice of statistic. The nksd posterior also falls within the general class of Gibbs posteriors, which have been studied in the context of robustness, randomized estimators, and generalized belief updating (Zhang, 2006a,b;Jiang and Tanner, 2008;Bissiri et al., 2016;Jewson et al., 2018;Miller and Dunson, 2019).
Our theoretical results also contribute to the emerging literature on Stein discrepancies (Anastasiou et al., 2021).  recently proposed minimum kernelized Stein discrepancy estimators and established their consistency and asymptotic normality. In Section 6, we establish a Bayesian counterpart to these results, showing that the nksd posterior is asymptotically normal (in the sense of Bernstein-von Mises) and admits a Laplace approximation. To prove this result, we rely on the recent work of Miller (2021) on the asymptotics of generalized posteriors. Since  show that the kernelized Stein discrepancy is related to the Hyvärinen divergence in that both are Stein discrepancies, our work bears an interesting relationship to that of Shao et al. (2018), who use a Bayesian version of the Hyvärinen divergence to perform model selection with improper priors. They derive a consistency result analogous to Equation 16, however, their model selection score takes the form of a prequential score, not a Gibbs marginal likelihood as in the SVC, and cannot be used for data selection.
In independent recent work, Matsubara et al. (2021) propose a Gibbs posterior based on the KSD and derive a Bernstein-von Mises theorem similar to Theorem 6.9 using the results of . Their method is not motivated by the Bayesian data selection problem but rather by (1) inference for energy-based models with intractable normalizing constants and (2) robustness to -contamination. Their Bernstein-von Mises theorem differs from ours in that it applies to a V-statistic estimator of the KSD rather than a U-statistic estimator of the NKSD.
Our linear approximation to the minimum Stein discrepancy estimator (Section 2.3.3) is directly inspired by the Swiss Army infinitesimal jackknife of Giordano et al. (2019), which similarly computes the linear response of an extremum estimator with respect to perturbations of the dataset.

Toy example
The purpose of this toy example is to illustrate the behavior of the Stein volume criterion, and compare it to some of the defective alternatives listed in Table 1, in a simple setting where all computations can be done analytically (Section S1.3). In all of the following experiments, we simulated data from a bivariate normal distribution: To set up the Stein volume criterion, we set T = 5 and we choose a radial basis function kernel, k(x, y) = exp(− 1 2 x − y 2 2 ), which factors across dimensions. We considered both dataset size-independent values of m B (in particular, m B = 5 r B ) and dataset size-dependent values of m B (in particular, Equation 3 with α = 0.5, θ = 1, and D = 0.2, where fractional values of D correspond to shared parameters across components in the Pitman-Yor mixture model), obtaining very similar results in each case (shown in Figures 2 and S1, respectively). These choices of m B ensure that, except for at very small N , the background model has more parameters per data dimension than each of the foreground models considered below, which have just one. In particular, m B > 1 r B for all N (in the size-independent case) and for N ≥ 5 (in the size-dependent case).

Data selection consistency
First, we set Σ 0 to be a diagonal matrix with entries (1, 1/2), that is, Σ 0 = diag(1, 1/2), and for x ∈ R 2 , we consider the model where I denotes the identity matrix. This parametric model is misspecified, owing to the incorrect choice of covariance matrix. We consider two choices of foreground subspace: the first dimension (defined by the projection matrix V F 1 = (1, 0) ) or the second dimension (projection matrix V F 2 = (0, 1) ). The model is only well-specified for F 1 (not F 2 ), so a successful data selection procedure would asymptotically select F 1 .
In Figure 2a, we see that the SVC correctly concentrates on F 1 as the number of datapoints N increases, with the log SVC ratio growing linearly in N , as predicted by Equation 16. Meanwhile, the naive alternative score K (a) (Equation 17) fails since it depends on the foreground entropies, while K (b) (Equation 21) succeeds since the volume correction is negligible in this case; see Section 3.1 and Table 1.

Nested data selection consistency
Next, we examine the nested data selection case. We use the same model (Equation 29), but we set Σ 0 = I so that the model is well-specified even without being projected. We compare the complete data space (X F 1 = X , projection matrix V F 1 = I) to the first dimension alone (projection matrix V F 1 = (1, 0) ). Nested data selection consistency demands that the higher-dimensional data space X F 1 be preferred asymptotically, since the model is wellspecified for both X F 1 and X F 2 . Figure 2b shows that this is indeed the case for the Stein volume criterion, with the log SVC ratio growing at a log N rate when m B is independent of N , as predicted by Equation 20. When m B depends on N via the Pitman-Yor expression, the log SVC ratio grows at a N α log N rate ( Figure S1b). Meanwhile, Figure 2b shows that K (a) and K (b) both fail to exhibit nested data selection consistency, in accordance with our theory (Section 3.2 and Table 1).

Model selection consistency (nested and non-nested)
Finally, we examine model selection and nested model selection consistency. We again set Σ 0 = I. We first compare the (well-specified) model q(x|θ) = N (x | θ, I) to the (misspecified) model q(x|θ) = N (x | θ, 2I), using the prior π(θ) = N (θ | (0, 0) , 10I) for both models. As shown in Figure 2c, the SVC correctly concentrates on the first model, with the log SVC ratio growing linearly in N , as predicted by Equation 25. The same asymptotic behavior is exhibited by K (a) , which is equivalent to the standard Bayesian marginal likelihood in this setting (Section 3.3). Finally, to check nested model selection consistency, we compare two well-specified nested models: q(x) = N (x | (0, 0) , I) and q(x|θ) = N (x | θ, I). Figure 2d shows that the SVC correctly selects the simpler model (that is, the model with smaller Here, we set m B = 5 r B . The plots show the results for 5 randomly generated datasets (thin lines) and the average over 100 random datasets (bold lines).
parameter dimension) and the log SVC ratio grows as log N (Equation 27). This, too, matches the behavior of the standard Bayesian marginal likelihood, seen in the plot of K (a) .
6 Theory 6.1 Properties of the NKSD Suppose X (1) , . . . , X (N ) are i.i.d. samples from a probability measure P on X ⊆ R d having density p(x) with respect to the Lebesgue measure. Let L 1 (P ) denote the set of measurable functions f such that f (x) p(x)dx < ∞ where · is the Euclidean norm. We impose the following regularity conditions to use the nksd to compare P with another probability measure Q having density q(x) with respect to the Lebesgue measure; these are similar to conditions used for the standard ksd in previous work .
Condition 6.1 (Restrictions on p and q). Assume s p (x) := ∇ x log p(x) and s q (x) := ∇ x log q(x) exist and are continuous for all x ∈ X , and assume X is connected and open. Further, assume s p , s q ∈ L 1 (P ).
We refer to s p as the Stein score function of p. Note that existence of s p (x) implies p(x) > 0. Now, consider a kernel k : X × X → R. The kernel k is said to be integrally strictly positive definite if for any g : Condition 6.2 (Restrictions on k). Assume the kernel k is symmetric, bounded, integrally strictly positive definite, and belongs to the Stein class of P .
The following result shows that the nksd can be written in a way that does not involve s p ; this is particularly useful for estimating the nksd when P is unknown. Proposition 6.3. If Conditions 6.1 and 6.2 hold, then the nksd is finite and The proof is in Section S3.1. Next, we show the nksd satisfies the properties of a divergence.
Proposition 6.4. If Conditions 6.1 and 6.2 hold, then with equality if and only if p(x) = q(x) almost everywhere.
The proof is in Section S3.1. Unlike the standard ksd, but like the kl divergence, the nksd exhibits subsystem independence (Caticha, 2004(Caticha, , 2011: if two distributions P and Q have the same independence structure, then the total nksd separates into a sum of individual nksd terms. This is formalized in Proposition 6.6. Condition 6.5 (Shared independence structure). Let x = (x 1 , x 2 ) be a decomposition of a vector x ∈ R d into two subvectors, x 1 and x 2 . Assume p(x) and q(x) factor as p(x) = p(x 1 )p(x 2 ) and q(x) = q(x 1 )q(x 2 ), and that the kernel k factors as k(x, y) = k 1 (x 1 , y 1 )k 2 (x 2 , y 2 ) where k 1 and k 2 both satisfy Condition 6.2.
Proposition 6.6 (Subsystem independence). If Conditions 6.1, 6.2, and 6.5 hold, then where the first term on the right-hand side uses kernel k 1 and the second term uses k 2 .
See Section S3.1 for the proof. Subsystem independence is powerful since it separates the problem of evaluating the foreground model from that of evaluating the background model. A modified version applies to the estimator nksd(p q) (Equation 5); see Proposition S3.1.

Bernstein-von Mises theorem for the NKSD posterior
In this section, we establish asymptotic properties of the SVC and, more broadly, of its corresponding generalized posterior, which we refer to as the nksd posterior, defined as In particular, in Theorem 6.9, we show that the nksd posterior concentrates and is asymptotically normal, and we establish that the Laplace approximation to the SVC (Equation 10) is asymptotically correct. These results form a Bayesian counterpart to those of , who establish the consistency and asymptotic normality of minimum ksd estimators. Thus, in both the frequentist and Bayesian contexts, we can replace the average log likelihood with the negative ksd and obtain similar key properties. Our results in this section do not depend on whether or not we are working with a foreground subspace, so we suppress the x F notation. Let Θ ⊆ R m , and let {Q θ : θ ∈ Θ} be a family of probability measures on X ⊆ R d having densities q θ (x) with respect to Lebesgue measure. For notational convenience, we sometimes write q(x|θ) instead of q θ (x). Suppose the data X (1) , . . . , X (N ) are i.i.d. samples from some probability measure P 0 on X having density p 0 (x) with respect to Lebesgue measure. To ensure the nksd satisfies the properties of a divergence for all q θ , and that convergence of nksd is uniform on compact subsets of Θ (Proposition S3.2), we require the following.
Condition 6.7. Assume Conditions 6.1 and 6.2 hold for p 0 , k, and q θ for all θ ∈ Θ. Further, assume that the kernel k has continuous and bounded partial derivatives up to and including second order, and k(x, y) > 0 for all x, y ∈ X . Now we can set up the generalized posterior. First define where u θ (x, y) is the u(x, y) function from Equation 5 with q θ in place of q. For the case of N = 1, we define f 1 (θ) = 0 by convention. Note that −N f N (θ) plays the role of the log likelihood. Also define where π(θ) is a prior density on Θ. Note that π N (θ)dθ is the NKSD posterior and z N is the corresponding generalized marginal likelihood employed in the SVC. Denote the gradient and Hessian of f by f (θ) = ∇ θ f (θ) and f (θ) = ∇ 2 θ f (θ), respectively. To ensure that the nksd posterior is well defined and has an isolated maximum, we assume the following condition.
Condition 6.8. Suppose Θ ⊆ R m is a convex set and (a) Θ is compact or (b) Θ is open and f N is convex on Θ with probability 1 for all N . Assume z N < ∞ a.s. for all N . Assume f has a unique minimizer θ * ∈ Θ, f (θ * ) is invertible, π is continuous at θ * , and π(θ * ) > 0.
By Proposition 6.4, f has a unique minimizer whenever {Q θ : θ ∈ Θ} is well-specified and identifiable, that is, when Q θ = P 0 for some θ and θ → Q θ is injective.
In Theorem 6.9 below, we establish the following results: (1) the minimum nksd converges to the minimum nksd; (2) π N concentrates around the minimizer of the nksd; (3) the Laplace approximation to z N is asymptotically correct; and (4) π N is asymptotically normal in the sense of Bernstein-von Mises. The primary regularity conditions we need for this theorem are restraints on the derivatives of s q θ with respect to θ (Condition 6.10). Our proof of Theorem 6.9 relies on the theory of generalized posteriors developed by . We use · for the Euclidean-Frobenius norms: Theorem 6.9. If Conditions 6.7, 6.8, and 6.10 hold, then there is a sequence θ N → θ * a.s. such that: 3.
The proof is in Section S3.2. We write ∇ 2 θ s q θ to denote the tensor in R d×m×m in which entry We write N to denote the set of natural numbers.
Condition 6.10 (Stein score regularity). Assume s q θ (x) has continuous third-order partial derivatives with respect to the entries of θ on Θ. Suppose that for any compact, convex subset C ⊆ Θ, there exist continuous functions g 0,C , g 1,C ∈ L 1 (P 0 ) such that for all θ ∈ C, x ∈ X , Further, assume there is an open, convex, bounded set E ⊆ Θ such that θ * ∈ E,Ē ⊆ Θ, and the sets are bounded with probability 1.
Next, Theorem 6.11 shows that in the special case where q θ (x) is an exponential family, many of the conditions of Theorem 6.9 are automatically satisfied.
Theorem 6.11. Suppose {Q θ : θ ∈ Θ} is an exponential family with densities of the form and assume Θ is convex, open, and nonempty. Assume log λ(x) and t(x) are continuously differentiable on X , ∇ x log λ(x) and ∇ x t(x) are in L 1 (P 0 ), and the rows of the Jacobian matrix ∇ x t(x) ∈ R m×d are linearly independent with positive probability under P 0 . Suppose Condition 6.7 holds, f has a unique minimizer θ * ∈ Θ, the prior π is continuous at θ * , and π(θ * ) > 0. Then the assumptions of Theorem 6.9 are satisfied for all N sufficiently large.
The proof is in Section S3.2.

Asymptotics of the Stein volume criterion
The Laplace approximation to the SVC uses the estimate nksd and its minimizer θ N , rather than the true nksd and its minimizer θ * . To establish the consistency properties of the SVC, we need to understand the relationship between the two. To do so, we adapt a standard approach to performing such an analysis of the marginal likelihood, for instance, as in Theorem 1 of Dawid (2011).
Theorem 6.12. Assume the conditions of Theorem 6.9 hold, and assume s q θ * and ∇ θ θ=θ * s q θ are in L 2 (P 0 ). Then as N → ∞, The proof is in Section S3.3. Remarkably, Equation 45 shows that f N (θ * ) converges to f (θ * ) more rapidly when the model is well-specified, specifically, at a 1/N rate instead of 1/ √ N . This is unusual and is crucial for our results in Section 6.4. The standard log likelihood does not exhibit this rapid convergence; see Section S2.1. This property of the nksd derives from similar properties exhibited by the standard ksd (Liu et al., 2016, Theorem 4.1). Combined with Theorem 6.9 (part 3), Theorem 6.12 implies that when the model is misspecified, the leading order term of log z N is −N f (θ * ), whereas when the model is well-specified, the leading order term is − 1 2 m log N .

Data and model selection consistency of the SVC
In this section, we establish the asymptotic consistency of the Stein volume criterion (SVC) when used for data selection, nested data selection, model selection, and nested model selection; see Theorem 6.17. This provides rigorous justification for the claims in Section 3. These results are all in the context of pairwise comparisons between two models or two model projections, M 1 and M 2 . Before proving the results, we formally define the consistency properties discussed in Section 3. Each property is defined in terms of a pairwise score ρ(M 1 , M 2 ), such as ρ(M 1 , M 2 ) = log(K 1 /K 2 ). For simplicity, we assume ρ(M 1 , M 2 ) = −ρ(M 2 , M 1 ); this is satisfied for all of the cases we consider. Let dim(·) denote the dimension of a real space.
Definition 6.13 (Data selection consistency). Consider foreground model projections Definition 6.14 (Nested data selection consistency). Consider foreground model projections Definition 6.16 (Nested model selection consistency). Consider foreground models M j := {q j (x F |θ j ) : θ j ∈ Θ j } for j ∈ {1, 2}. We say that ρ satisfies "nested model selection con- In each case, ρ may diverge almost surely ("strong consistency") or in probability ("weak consistency"). Note that in Definitions 6.13-6.14, the difference between M 1 and M 2 is the choice of foreground data space F, whereas in Definitions 6.15-6.16, M 1 and M 2 are over the same foreground space but employ different model spaces.
In Theorem 6.17, we show that the SVC has the asymptotic properties outlined in Section 3. In combination with the subsystem independence properties of the NKSD (Propositions 6.6 and S3.1), Theorem 6.17 also leads to the conclusion that the SVC approximates the NKSD marginal likelihood of the augmented model (Equation 8). Our proof is similar in spirit to previous results for model selection with the standard marginal likelihood, notably those of  and Huggins and Miller (2021), but relies on the special properties of the nksd marginal likelihood in Theorem 6.12.
Theorem 6.17. For j ∈ {1, 2}, assume the conditions of Theorem 6.12 hold for model M j defined on X F j , with density q j (x F j |θ j ) for θ j ∈ Θ j ⊆ R m F j ,j . Let K j,N be the Stein volume criterion for M j , with background model penalty m B j = m B j (N ), and let θ j, * := argmin θ j nksd(p 0 (x F j ) q j (x F j |θ j )). Then: where c B 1 and c B 2 are positive and constant in N , then The proof is in Section S3.4. In particular, assuming the conditions of Theorem 6.12, we obtain the following consistency results in terms of convergence in probability.
• If m B j = o(N/ log N ) then the SVC exhibits data selection consistency and model selection consistency. This holds by Theorem 6.17 (part 1) since D 2 > D 1 = 0.
• Consider a nested data selection problem with N and c B 2 > c B 1 > 0, then the SVC exhibits nested data selection consistency. Cases A and B hold by Theorem 6.17 (parts 2 and 3, respectively) since D 1 = D 2 = 0.
7 Application: Probabilistic PCA Probabilistic principal components analysis (pPCA) is a commonly used tool for modeling and visualization. The basic idea is to model the data as linear combinations of k latent factors plus Gaussian noise. The inferred weights on the factors are frequently used to provide low-dimensional summaries of the data, while the factors themselves describe major axes of variation in the data. In practice, pPCA is often applied in settings where it is likely to be misspecified -for instance, the weights are often clearly non-Gaussian. In this section, we show how data selection can be used to uncover sources of misspecification and to analyze how this misspecification affects downstream inferences.
The generative model used in pPCA is independently for i = 1, . . . , N , where I k is the k-dimensional identity matrix, Z (i) ∈ R k is the weight vector for datapoint i, H ∈ R d×k is the unknown matrix of latent factors, and v > 0 is the variance of the noise. To form a Laplace approximation for the Stein volume criterion, we follow the approach developed by  for the standard marginal likelihood. Specifically, we parameterize H as where U is a d × k matrix with orthonormal columns (that is, it lies on the Stiefel manifold) and L is a k × k diagonal matrix. We use the priors suggested by Minka (2001), where U is the set of d × k matrices with orthonormal columns and L ii is the ith diagonal entry of L. We set α = 0.1 in the following experiments, and we use pymanopt  to optimize U over the Stiefel manifold (Section S4).

Simulations
In simulations, we evaluate the ability of the SVC to detect partial misspecification. We set d = 6, draw the first four dimensions from a pPCA model with k = 2 and and generate dimensions 5 and 6 in such a way that pPCA is misspecified. We consider two misspecified scenarios: scenario A (Figure 3a) is that where Σ W (i) = (0.05) W (i) I 2 . Scenario B (Figure 3d) is the same but with Scenario B is more challenging because the marginals of the misspecified dimensions are still Gaussian, and thus, misspecification only comes from the dependence between X 5 and X 6 . As illustrated in Figures 3b and 3e, both kinds of misspecification are very hard to see in the lower-dimensional latent representation of the data. Our method can be used to both (i) detect misspecified subsets of dimensions, and (ii) conversely, find a maximal subset of dimensions for which the pPCA model provides a  We performed leave-one-out data selection, comparing the foreground space X F 0 = X to foreground spaces X F j for j ∈ {1, . . . , d}, which exclude the jth dimension of the data. We computed the log SVC ratio log(K j /K 0 ) = log K j −log K 0 using the BIC approximation to the SVC (Section 2.3.2) and the approximate optima technique (Section 2.3.3). We quantify the performance of the method in detecting misspecified dimensions in terms of the balanced accuracy, defined as (T N/N + T P/P )/2 where T N is the number of true negatives (dimension by dimension), N is the number of negatives, T P is the number of true positives, and P is the number of positives. Experiments were repeated independently five times. Figures 3c and 3f show that as the sample size increases, the SVC correctly infers that dimensions 1 through 4 should be included and dimensions 5 and 6 should be excluded.

Comparison with a nonparametric background model
To benchmark our method, we compare with an alternative approach that uses an explicit augmented model. The Pólya tree is a nonparametric model with a closed-form marginal likelihood that is tractable for one-dimensional data . We define a flexible background model by sampling each dimension j of the background space independently as with the Pólya tree constructed as by  (Section S4.4). We set F = N (0, 10),F = N (0, 10), and η = 1000 so that the model is weighted only very weakly towards the base distribution. We performed data selection using the marginal likelihood of the Pólya tree augmented model, computing the marginal of the pPCA foreground model using the approximation of . The accuracy results for data selection are in Figures 3c and 3f. On scenario A (Equation 50), the Pólya tree augmented model requires significantly more data to detect which dimensions are misspecified. On scenario B (Equation 51) the Pólya tree augmented model fails entirely, preferring the full data space X F 0 = X which includes all dimensions (Figure 3f). The reason is that the background model is misspecified due to the assumption of independent dimensions, and thus, the asymptotic data selection results (Equations 15 and 19) do not hold. This could be resolved by using a richer background model that allows for dependence between dimensions, however, computing the marginal likelihood under such a model would be computationally challenging. Even with the independence assumption, the Pólya tree approach is already substantially slower than the SVC (Figure 3g).

Application to pPCA for single-cell RNA sequencing
Single-cell RNA sequencing (scRNAseq) has emerged as a powerful technology for highthroughput characterization of individual cells. It provides a snapshot of the transcriptional state of each cell by measuring the number of RNA transcripts from each gene. PCA is widely used to study scRNAseq datasets, both as a method for visualizing different cell types in the dataset and as a pre-processing technique, where the latent embedding is used for downstream tasks like clustering and lineage reconstruction (Qiu et al., 2017;van Dijk et al., 2018). We applied data selection to answer two practical questions in the application of probabilistic PCA to scRNAseq data: (1) Where is the pPCA model misspecified? (2) How does partial misspecification of the pPCA model affect downstream inferences?

Model criticism
Our first goal was to verify that the SVC provides reasonable inferences of partial model misspecification in practice. We examined two different scRNAseq datasets, focusing for illustration on a dataset from human peripheral blood mononuclear cells taken from a healthy donor, and pre-processed the data following standard procedures in the field (Section S4.5). We subsampled each dataset to 200 genes (selected randomly from among the 2000 most highly expressed) and 2000 cells (selected randomly) for computational tractability, then mean-subtracted and standardized the variance of each gene, again following standard practice in the field. The number of latent components k was set to 3, based on the procedure of Minka (2000). We performed leave-one-out data selection, comparing the foreground space X F 0 := X to foreground spaces X F j that exclude the jth gene. We computed the log SVC ratio log K j − log K 0 using the BIC approximation to the SVC (Section 2.3.2) and the approximate optima technique (Section 2.3.3). We used the same setting of T and of m B as was used in simulation, resulting in a background model complexity of m B = 20 r B for datasets of this size. Based on the SVC criterion, 162 out of 200 genes should be excluded from the foreground pPCA model, suggesting widespread partial misspecification. Figure 4 compares the histogram of individual genes to their estimated density under the pPCA model inferred for X F 0 = X . Those genes most favored to be excluded (namely, UBE2V2 and IRF8) show extreme violations of normality, in stark contrast to those genes most favored to be included (MT-CO1 and RPL6).
Next, we compared the results of our data selection approach to a more conventional strategy for model criticism. Criticism of partially misspecified models can be challenging in practice because misspecification of the model over some dimensions of the data can lead to substantial model-data mismatch in dimensions for which the model is indeed wellspecified (Jacob et al., 2017). The standard approach to model criticism-first fit a model, then identify aspects of the data that the model poorly explains-can therefore be misleading if our aim is to determine how the model might be improved (e.g., in the context of "Box's loop", Blei, 2014). In particular, standard approaches such as posterior predictive checks will be expected to overstate problems with components of the model that are well-specified and understate problems with components of the model that are misspecified. Bayesian data selection circumvents this issue by evaluating augmented models, which replace potentially misspecified components of the model by well-specified components. To illustrate the difference between these approaches in practice, we compared the SVC to a closely analogous measurement of error for the full foreground model (inferred from X F 0 = X ), where θ 0,N := argmin nksd(p 0 (x) q(x|θ)) is the minimum nksd estimator for the foreground model when including all dimensions. This model criticism score evaluates the amount of model-data mismatch contributed by the subspace X B j when modeling all data dimensions with the foreground model. For comparison, the BIC approximation to the log SVC ratio is where θ j,N := argmin nksd(p 0 (x F j ) q(x F j |θ)) is the minimum nksd estimator for the projected foreground model applied to the restricted dataset, which we approximate as θ 0,N plus the implicit function correction derived in Section 2.3.3. Figure 5 illustrates the differences between the conventional criticism approach (log E j − log E 0 ) and the log SVC ratio on an scRNAseq dataset. To enable direct comparison of the two methods, we focus on the lower order terms of Equation 54, that is, we set m B j = m F 0 − m F j . We see that the amount of error contributed by X B j , as judged by the SVC, is often substantially higher than the amount indicated by the conventional criticism approach, implying that the conventional criticism approach understates the problems caused by individual genes and, conversely, overstates , and the conventional full model criticism score, log E j − log E 0 , for each gene.
the problems with the rest of the model. Using the SVC instead of a standard criticism approach can also help clarify trends in where the proposed model fails. A prominent concern in scRNAseq data analysis is the common occurrence of cells that show exactly zero expression of a certain gene (Pierson and Yau, 2015;Hicks et al., 2018). We found a Spearman correlation of ρ = 0.89 between the conventional criticism log E j − log E 0 for a gene j and the fraction of cells with zero expression of that gene j, suggesting that this is an important source of model-data mismatch in this scR-NAseq dataset, but not necessarily the only source ( Figure 6a). However, the log SVC ratio yields a Spearman correlation of ρ = 0.98, suggesting instead that the amount of model-data mismatch can be entirely explained by the fraction of cells with zero expression (Figure 6b). These observations are repeatable across different scRNAseq datasets (Figure 6c, 6d).

Evaluating robustness
Data selection can also be used to evaluate the robustness of the foreground model to partial model misspecification. This is particularly relevant for pPCA on scRNAseq data, since the inferred latent embeddings of each cell are often used for downstream tasks such as clustering, lineage reconstruction, and so on. Misspecification may produce spurious conclusions, or alternatively, misspecification may be due to structure in the data that is scientifically interesting. To understand how partial misspecification of the pPCA model affects the latent representation of cells (and thus, downstream inferences), we performed data selection with a sequence of background model complexities c B , where m B = c B r B (Figure 7a). We inferred the pPCA parameters based only on genes that the SVC selects to include in the foreground subspace. Figures 7e-7b visualize how the latent representation changes as c B grows and fewer genes are selected. We can observe the representation morphing into a standard nor- mal distribution, as we would expect in the case where the pPCA model is well-specified. However, the relative spatial organization of cells in the latent space remains fairly stable, suggesting that this aspect of the latent embedding is robust to partial misspecification. We can conclude that, at least in this example, misspecification strongly contributes to the non-Gaussian shape of the latent representation of the dataset, but not to the distinction between subpopulations. This suggests a simple physical analogy: if individual gene expression is a two-state system, we might study gene regulation with the theory of interacting two-state systems, namely spin glasses. We can consider for instance a standard model of this type in which each cell i is described by a vector of spins z i = (z i1 , . . . , z id ) drawn from an Ising model, specifying whether each gene j ∈ {1, . . . , d} is "on" or "off". In reality, gene expression lies on a continuum, so we use a continuous relaxation of the Ising model and parameterize each spin using a logistic function, setting z ij1 (x ij , µ, τ ) = 1/(1 + exp(−τ (x ij − µ))) and z ij2 (x ij , µ, τ ) = 1 − z ij1 (x ij , µ, τ ). Here, x ij is the observed expression level of gene j in cell i, the unknown parameter µ controls the threshold for whether the expression of a gene is "on" (such that z ij ≈ (1, 0) ) or "off" (such that z ij ≈ (0, 1) ), and the unknown parameter τ > 0 controls the sharpness of the threshold. The complete model is then given by where Z H,J,µ,τ is the unknown normalizing constant of the model, and the vectors H j ∈ R 2 and matrices J jj ∈ R 2×2 are unknown parameters. This model is motivated by experimental observations and is closely related to RNAseq analysis methods that have been successfully applied in the past (Friedman et al., 2000;Friedman, 2004;Ding and Peng, 2005;Chen et al., 2015;Banerjee et al., 2008;Duvenaud et al., 2008;Liu et al., 2009;Huynh-Thu et al., 2010;Moignard et al., 2015;Matsumoto et al., 2017). However, from a biological perspective we can expect that serious problems may occur when applying the model naively to an scR-NAseq dataset. Genes need not exhibit bistable expression: it is straightforward in theory to write down models of gene regulation that do not have just one or two steady states-gene expression may fall on a continuum, or oscillate, or have three stable states-and many alternative patterns have been well-documented empirically (Alon, 2019). Interactions between genes may also be more complex than the model assumes, involving for instance three-way dependencies between genes. All of these biological concerns can potentially produce severe violations of the proposed two-state glass model's assumptions. Data selection provides a method for discovering where the proposed model applies.
Applying standard Bayesian inference to the glass model is intractable, since the normalizing constant is unknown (it is an energy-based model). However, the normalizing constant does not affect the SVC, so we can still perform data selection. We used a variational approximation to the SVC (Section 2.3.4). We placed a Gaussian prior on H and a Laplace prior on each entry of J to encourage sparsity in the pairwise gene interactions; we also used Gaussian priors for µ and τ after applying an appropriate transform to remove constraints (Section S5.1). Following the logic of stochastic variational inference, we optimized the variational approximation using minibatches of the data and a reparameterization gradient estimator (Hoffman et al., 2013;Kingma and Welling, 2014;Kucukelbir et al., 2017). We also simultaneously stochastically optimized the set of genes included in the foreground subspace, using the Leave-One-Out REINFORCE estimator (Kool et al., 2019;Dimitriev and Zhou, 2021). We implemented the model and inference strategy within the probabilistic programming language Pyro by defining a new distribution with log probability given by the negative NKSD . Pyro provides automated, GPU-accelerated stochastic variational inference, requiring less than an hour for inference on datasets with thousands of cells.
We examined three scRNAseq datasets, taken from (i) peripheral blood monocytes (PBMCs) from a healthy donor (2,428 cells), (ii) a MALT lymphoma (7,570 cells), and (iii) mouse neurons (10,658 cells) (Section S5.2). We preprocessed the data following standard protocols and focused on 200 high expression, high variability genes in each dataset, based on the metric of Gigante et al. (2020). We set T = 0.05 as in Section 7, and used the Pitman-Yor expression for m B (Equation 3) with α = 0.5, θ = 1 and D = 100. This value of D ensures that the number of background model parameters per data dimension is larger than the number of foreground model parameters per data dimension except for at very small N ; in particular, there are 798 foreground model parameter dimensions associated with each
We investigated the biological information captured by the foreground model on the MALT dataset. In particular, we looked at the approximate NKSD posterior for the selected 187 genes, and compared it to the approximate NKSD posterior for the model when applied to all 200 genes. (Note that, since the glass model lacks a tractable normalizing constant, we cannot compare standard Bayesian posteriors.) Figure 8 shows, for a subset of selected genes, the posterior mean of the interaction energy ∆E jj := J jj 21 + J jj 12 − J jj 22 − J jj 11 , that is, the total difference in energy between two genes being in the same state versus in opposite states. We focused on strong interactions with |∆E jj | > 1, corresponding to just 5% of all possible gene-gene interactions ( Figure S3).
One foreground gene with especially large loading onto the top principal component of the ∆E matrix is CD37 (Figure 8). In B-cell lymphomas, of which MALT lymphoma is an example, CD37 loss is known to be associated with decreased patient survival (Xu-Monette et al., 2016). Further, previous studies have observed that CD37 loss leads to high NF-κB pathway activation (Xu-Monette et al., 2016). Consistent with this observation, the estimated interaction energies in our model suggest that decreasing CD37 will lead to higher expression of REL, an NF-κB transcription factor (∆E CD37,REL = 2.5), decreased expression of NKFBIA, an NF-κB inhibitor (∆E CD37,NKFBIA = −3.6), and higher expression of BCL2A1, a downstream target of the NF-κB pathway (∆E CD37,BCL2A1 = 2.1). Separately, a knockout study of Cd37 in B-cell lymphoma in mice does not show IgM expression (de Winde et al., 2016), consistent with our model (∆E CD37,IGHM = −8.2). The same study does show MHC-II expression, and our model predicts the same result, for HLA-DQ in particular (∆E CD37,HLA-DQA1 = 5.0, ∆E CD37,HLA-DQB1 = 3.7). These results suggest that the data selection procedure can successfully find systems of interacting genes that can plausibly be modeled as a spin glass, and which, in this case, are relevant for cancer.
To investigate whether data selection provided a benefit in this analysis, we compare with the results obtained by applying the foreground model to the full dataset of all 200 genes. All but one of the interactions listed above have |∆E| < 1 in the full foreground model, and three have opposite signs (∆E CD37,NFKBIA = +0.7, ∆E CD37,IGHM = +0.0, ∆E CD37,HLA-DQB1 = −0.6); see Figure S4. Across all 187 selected genes, we find only a moderate correlation between the interaction energies estimated when using the full foreground model compared with the data selection-based model (Spearman's rho = 0.30, p < 0.01; Figure 9). These results show that using data selection can lead to substantially different, and arguably more biologically plausible, downstream conclusions as compared to naive application of the foreground model to the full dataset.
As a simple alternative, one might wonder whether genes that are poorly fit by the model could be identified simply by looking their posterior uncertainty under the full foreground model. This simple approach does not work well, however, since it is possible for parameters to have low uncertainty even when the model poorly describes the data. Indeed, we found that examining uncertainty in the glass model does not lead to the same conclusions as performing data selection: the genes excluded by our data selection procedure are not the ones with the highest uncertainty in their interactions (as measured by the mean posterior standard deviation of ∆E jj under the NKSD posterior), though they do have above average uncertainty ( Figure S5a). Instead, the genes excluded by our data selection procedure are the ones with the highest fraction of cells with zero expression, violating the assumptions of the foreground model ( Figure S5b). These results show how data selection provides a sound, computationally tractable approach to criticizing and evaluating complex Bayesian models.

Discussion
Statistical modeling is often described as an iterative process, where we design models, infer hidden parameters, critique model performance, and then use what we have learned from the critique to design new models and repeat the process (Gelman et al., 2013). This process has been called "Box's loop" (Blei, 2014). From one perspective, data selection offers a new criticism approach. It goes beyond posterior predictive checks and related methods by changing the model itself, replacing potentially misspecified components with a flexible background model. This has important practical consequences: since misspecification can distort estimates of model parameters in unpredictable ways, predictive checks are likely to indicate mismatch between the model and the data across the entire space X even when the proposed parametric model is only partially misspecified. Our method, by contrast, reveals precisely those subspaces of X where model-data mismatch occurs. From another perspective, data selection is outside the design-infer-critique loop. An underlying assumption of Box's loop is that scientists want to model the entire dataset. As datasets get larger, and measurements get more extensive, this desire has led to more and more complex (and often difficult to interpret) models. In experimental science, however, scientists have often followed the opposite trajectory: faced with a complicated natural phenomenon, they attempt to isolate a simpler example of the phenomenon for close study. Data selection offers one approach to formalizing this intuitive idea in the context of statistical analysis: we can propose a simple parametric model and then isolate a piece of the whole dataset-a subspace X F -to which this model applies. When working with large, complicated datasets, this provides a method of searching for simpler phenomena that are hypothesized to exist.
Supplementary material for "Bayesian data selection" S1 Methods details S1.1 Calibrating T The SVC contains a hyperparameter T > 0. To choose an appropriate value of T , we aim, roughly, to match the coverage of the generalized posterior to the coverage of the standard Bayesian posterior when the model is well-specified. Let θ * be the true parameter value, such that p 0 (x) = q(x|θ * ) almost everywhere. Let Under regularity conditions , according to the Bernstein-von Mises theorem, h kl N converges to a normal distribution in total variation, According to Theorem 6.9, the generalized posterior associated with the SVC has analogous behavior. Let G svc (θ) := ∇ 2 θ 1 T nksd(p 0 (x) q(x|θ)) and let θ svc N := argmin nksd(p 0 (x) q(x|θ)). Let h svc N be the density of √ N (θ −θ svc N ) when θ ∼ π svc N . Then by Theorem 6.9, h svc N converges to a normal distribution in total variation, For the uncertainty in each posterior to be roughly the same order of magnitude, we want or equivalently,

S1
To choose a single T value, we simulate true parameters from the prior, generate data from each simulated true parameter, and take the median of the estimated T values. That is, we use the medianT across samples drawn as In practice, we find that the order of magnitude ofT is stable across samples θ * from π(θ). See Section S4.3 for an example.

S1.2 Kernel recommendations
To obtain subsystem independence (Proposition 6.6), we suggest using a kernel that factors across subspaces, such that k(X, where k F and k B are integrally strictly positive definite kernels. In the applications in Sections 7 and 8, we use the following kernel. Definition S1.1. The factored inverse multiquadric (IMQ) kernel is defined as Note that this kernel factors across any subset of dimensions, that is, if S ⊆ {1, . . . , d} and S c = {1, . . . , d} \ S, then we can write k(x, y) = k S (x S , y S )k S c (x S c , y S c ). Thus, if the foreground subspace X F is the span of a subset of the standard basis, such that x F = V x = x S for some S ⊆ {1, . . . , d}, then it follows that k factors as k(x, y) = k F (x F , y F )k B (x B , y B ). Along with this observation, the next result shows that the factored IMQ satisfies the conditions of Theorem 6.9 that pertain to k alone.
Proposition S1.2. The factored IMQ kernel is symmetric, positive, bounded, integrally strictly positive definite, and has continuous and bounded partial derivatives up to order 2.
Proof. It is clear that k(x, y) = k(y, x) and k(x, y) > 0. Next, we show that k has continuous and bounded partial derivatives up to order 2. Note that we can write k(x, y) Since r 2 ≥ 0 and β < 0, |ψ(r)| ≤ c 2β/d for all r ∈ R. Further, it is straightforward to verify that |ψ (r)| and |ψ (r)| are bounded on R by using the fact that |r|/(c 2 + r 2 ) ≤ 1/(2c). By the chain rule, it follows that for all i, j, the functions k(x, y), |∂k/∂x i |, and |∂ 2 k/∂x i ∂y j | are bounded. Thus, we conclude that k, ∇k , and ∇ 2 k are bounded.
Finally, we show that k is integrally strictly positive definite. First, for any d, for x, y ∈ R d , the function (x, y) → (c 2 + x − y 2 2 ) β/d is an integrally strictly positive definite kernel (see, for example, Section 3.1 of Sriperumbudur et al., 2010); we refer to this as the standard IMQ kernel. Since the factored IMQ is a product of one-dimensional standard IMQ kernels, it defines a kernel on R d (Lemma 4.6 of Steinwart and Christmann, 2008) and is positive definite (Theorem 4.16 of Steinwart and Christmann, 2008). By Bochner's theorem (Theorem 3 of Sriperumbudur et al., 2010), a continuous positive definite kernel can be expressed in terms of the Fourier transform of a finite nonnegative Borel measure. In particular, applying Bochner's theorem to ψ(r), we have by Fubini's theorem, where Λ 0 is the finite nonnegative Borel measure on R associated with ψ(r) and Λ = Λ 0 × · · · × Λ 0 is the resulting product measure on R d . Applying Bochner's theorem in the other direction, we see that Ψ is a positive definite function. Moreover, since the standard IMQ kernel is characteristic (Theorem 7 of Sriperumbudur et al., 2010), it follows that the support of Λ 0 is R (Theorem 9 of Sriperumbudur et al., 2010), and thus the support of Λ is R d . This implies that the factored IMQ kernel k is characteristic (Theorem 9 of Sriperumbudur et al., 2010) and, since k is also translation invariant, k must be integrally strictly positive definite (Section 3.4 of Sriperumbudur et al., 2011).
Our choice of the factored IMQ kernel is motivated by the analysis of Gorham and Mackey (2017), which suggests that the standard IMQ is a good default choice for the kernelized Stein discrepancy, particularly when working with distributions that are (roughly speaking) very spread out. In particular, it is straightforward to show that the factored IMQ kernel, like the standard IMQ kernel, meets the conditions of Theorem 3.2 of Huggins and Mackey (2018). However, we do not pursue further the question of whether the nksd with the factored IMQ detects convergence and non-convergence since our statistical setting is different from that of Gorham and Mackey (2017), and we are assuming the data consists of i.i.d. samples from some underlying distribution rather than correlated samples from an MCMC chain which may or may not converge.

S2 Asymptotics of the alternative selection criteria
Theorem 6.17 shows that the SVC exhibits all four types of consistency: data selection, nested data selection, model selection, and nested model selection. In this section, we establish the consistency properties of the alternative criteria considered in Section 3.

S2.1 Setup
We first review the asymptotics of the standard marginal likelihood, discussed in depth by Dawid (2011) and , for example. Define Let m be the dimension of the parameter space. Under suitable regularity conditions , the Laplace approximation to the marginal likelihood is almost surely, where a N ∼ b N indicates that a N /b N → 1 as N → ∞. We can rewrite this as As shown by Dawid (2011) and , under regularity conditions, The nksd marginal likelihood has a similar decomposition. Following Section 6, define As shown in Theorem 6.9, almost surely as N → ∞. As above, we can rewrite this as By Theorem 6.12, we have and further, when the model is well-specified, such that nksd(p 0 (x) q(x|θ nksd * )) = 0, For ease of reference, here are the various scores that we consider for model/data selection. Marginal likelihood of the augmented model (foreground+background): Foreground marginal nksd, background volume correction (a.k.a. the SVC): Foreground marginal likelihood, background volume correction: ).

S2.2 Data selection
Assume m B j = o(N/ log N ) for j ∈ {1, 2}. By Equations S5 and S6, so K (a) does not satisfy data selection consistency. The SVC satisfies data selection consistency by Theorem 6.17 (part 1). We show that the other scores also satisfy data selection consistency. Since K (b) = (2π/N ) −m B /2 K where K is the SVC, by Theorem 6.17 (part 1), By Equation S10 and the fact that K (c) = exp(N H F )K (a) , we have These methods therefore satisfy data selection consistency. For the marginal likelihood of the augmented model, suppose m B 1 and m B 2 do not depend on N . Then by Equation S5, We can rewrite this in terms of the KL divergence. First note the decomposition, for j ∈ {1, 2}. Adding and subtracting the entropy H in Equation S15, and using the fact that the background model is well-specified,

S2.3 Nested data selection
In nested data selection, we are concerned with situations in which X F 2 ⊂ X F 1 and the model is well-specified over both X F 1 and X F 2 . Assume further that m B 2 − m B 1 does not depend on N . First, consider K (d) and K BIC . Since ) and by Theorem 6.12, f nksd Likewise, since K BIC = (2π/N ) m F /2 K (d) , it follows that

S9
As in Section 6.4, it is natural to assume m B 2 > m B 1 and m F 2 + m B 2 > m F 1 + m B 1 , in which case these criteria satisfy nested data selection consistency. None of K (a) , K (b) , and K (c) are guaranteed to satisfy nested data selection consistency, because the contribution of background model complexity is negligible or nonexistent. To see this, note that assuming m B j = o(N/ log N ), by Equation S10 we have Meanwhile, since K (b) = (2π/N ) −m B /2 K then by Theorem 6.17 (part 2), Since X F 2 ⊂ X F 1 , we have m F 2 ≤ m F 1 except perhaps in highly contrived scenarios. If m F 2 < m F 1 then Equation S20 shows that log(K On the other hand, if m F 2 = m F 1 , then by Equations S7 and S8, log(K ), then by Equations S5 and S6, If σ 2 := V P 0 (log p 0 (X F 1 )/p 0 (X F 2 )) is positive and finite, then by the central limit theorem and Slutsky's theorem, N −1/2 log(K (c) 2 ) D − → N (0, σ 2 ). Thus, K (c) randomly selects F 1 or F 2 with equal probability, and therefore, it does not satisfy nested data selection consistency.
For the marginal likelihood of the augmented model, suppose m B 1 and m B 2 do not depend on N . The marginal likelihood achieves nested data selection consistency because the augmented models are both well-specified and describe the complete data space X ; this guarantees that the O P 0 ( √ N ) terms in the marginal likelihood decomposition cancel. Specifically, p 0 (x) = q(x | θ kl j, * , φ kl j, * , F j ) for j ∈ {1, 2}, and thus, by Equations S5 and S6 applied to the augmented model, Nested data selection consistency follows assuming m F 2 + m B 2 > m F 1 + m B 1 as before. This can be contrasted with Equation S21, where although both foreground models are wellspecified, they describe different data (X

S2.4 Model selection
All of the criteria we consider satisfy model selection consistency. To see this, we apply the same asymptotic analysis as used for data selection in Section S2.2, under the same conditions on m B , obtaining 1 N logq 1 (X (1:N ) |F) q 2 (X (1:N ) |F) Note that in contrast to the data selection case, K (a) satisfies model selection consistency since the entropy terms H F j cancel due to the fact that F is fixed. We can think of this as a consequence of the kl divergence's subsystem independence; if we are just interested in modeling a fixed foreground space, there is no problem considering the foreground marginal likelihood alone (Caticha, 2004(Caticha, , 2011.

S2.5 Nested model selection
In nested model selection, since both models are well-specified, we have q j (x F |θ kl j, * ) = p 0 (x F ) = q j (x F |θ nksd j, * ) for j ∈ {1, 2}. Thus, the estimated divergences cancel: Using this along with Equations S5-S9, under the same conditions on m B as in Section S2.2, where we are using the assumption that the background model is the same in the two augmented modelsq 1 andq 2 and so m B,1 = m B,2 . Only K (d) fails to satisfy nested model selection consistency.

S3.1 Proofs of NKSD properties
Proof of Proposition 6.3. By assumption, the kernel is bounded, say |k(x, y)| ≤ B, and s p , s q ∈ L 1 (P ). Thus, by the Cauchy-Schwarz inequality, Since the kernel is integrally strictly positive definite and |k(x, y)| ≤ B, Thus, the nksd is finite. Equation 30 follows from Theorem 3.6 of .
Proof of Proposition 6.4. The denominator of the nksd is positive since k is integrally strictly positive definite. Defining δ(x) = s q (x) − s p (x), the numerator of the nksd is If δ i (x)p(x) = 0 almost everywhere with respect to Lebesgue measure on X , then the ith term on the right-hand side is zero. Meanwhile, if δ i (x)p(x) is not a.e. zero, then X |δ i (x)|p(x)dx > 0, and hence, the ith term is positive since k is integrally strictly positive definite and δ i ∈ L 1 (P ) by assumption. Hence, the nksd is nonnegative, and equals zero if and only if δ(x)p(x) = 0 almost everywhere. Suppose δ(x)p(x) = 0 almost everywhere. Since p(x) > 0 on X by assumption, this implies s p (x) = s q (x) a.e., and in fact, s p (x) = s q (x) for all x ∈ X by continuity. Since X is open and connected, then by the gradient theorem (that is, the fundamental theorem of calculus for line integrals), p(x) ∝ q(x), and hence, p(x) = q(x) on X . Conversely, if p(x) = q(x) almost everywhere, then δ(x)p(x) = 0 almost everywhere.
The result follows.
S3.2 Proof of Theorems 6.9 and 6.11 Our proofs in this section build on the proof of Theorem 3 of .
Proposition S3.2. Under the assumptions of Theorem 6.9, for any compact convex C ⊆ Θ, Proof. First, we establish almost sure convergence for the denominator of f N (θ). Since k is assumed to be bounded and to have bounded derivatives up to order two, we can choose B < ∞ such that B ≥ |k| + ∇ x k + ∇ x ∇ y k . In particular, the expected value of the kernel is finite: By the strong law of large numbers for U-statistics (Theorem 5.4A of Serfling, 2009), Note that the limit is positive since k(x, y) > 0 for all x, y ∈ X . For the numerator, we establish bounds on u θ and ∇ θ u θ . Let C ⊆ Θ be compact and convex. By Equation 5, for all θ ∈ C and all x, y ∈ X , |u θ (x, y)| ≤ |s q θ (x) s q θ (y)k(x, y)| + |s q θ (x) ∇ y k(x, y)| Similarly, for all θ ∈ C and all x, y ∈ X , ≤ g 0,C (x)g 1,C (y)B + g 0,C (y)g 1,C (x)B + g 1,C (x)B + g 1,C (y)B =: h 1,C (x, y).

(S42)
Note that h 0,C and h 1,C are continuous and belong to L 1 (P 0 × P 0 ). Let S 1 ⊆ S 2 ⊆ · · · ⊆ X be a sequence of compact sets such that ∪ ∞ M =1 S M = X . Note that this implies ∪ ∞ M =1 S M × S M = X × X . Suppose for the moment that, for each M , the following collections of functions are equicontinuous on C: (A) (θ → u θ (x, y) : x, y ∈ S M ) and (B) θ → u θ (x, y)P 0 (dy) : x ∈ S M . Assuming this, Theorem 1 of Yeo and Johnson (2001) shows that and that θ → X X u θ (x, y)P 0 (dx)P 0 (dy) is continuous. (Note that although Yeo and Johnson (2001) assume X = R, their proof goes through without further modification for any nonempty X ⊆ R d .) Combining Equations S40 and S43, we have Thus, it follows that sup θ∈C |f N (θ) − f (θ)| → 0 a.s. by Equations S40 and S41. To complete the proof, we must show that (A) and (B) are equicontinuous on C.
(A) Since θ → u θ (x, y) is differentiable on C, then by the mean value theorem, we have that for all θ 1 , θ 2 ∈ C and all x, y ∈ S M , Here, the second inequality holds sinceθ ∈ C by the convexity of C, and the supremum is finite because a continuous function on a compact set attains its maximum. Therefore, (θ → u θ (x, y) : x, y ∈ S M ) is equicontinuous on C.
Further, due to Equations S41 and S42, we can apply the Leibniz integral rule (Folland, 1999, Theorem 2.27) and find that ∇ θ u θ (x, y)P 0 (dy) exists and is equal to ∇ θ u θ (x, y)P 0 (dy). Now we apply the mean value theorem and the same reasoning as before to find that for all θ 1 , θ 2 ∈ C and all x ∈ S M , h 1,C (x, y)P 0 (dy) < ∞ whereθ = γθ 1 + (1 − γ)θ 2 for some γ ∈ [0, 1]. The supremum is finite since x → h 1,C (x, y)P 0 (dy) is continuous, which can easily be seen by plugging in the definition of h 1,C . Therefore, θ → u θ (x, y)P 0 (dy) : x ∈ S M is equicontinuous on C. Proposition S3.3. Under the assumptions of Theorem 6.9, (f N : N ∈ N) is uniformly bounded on E.
Proof. First, for any x, y ∈ X , if we define g(θ) = s q θ (x) and h(θ) = s q θ (y) then u θ = (g h)k + g (∇ y k) + h (∇ x k) + trace(∇ x ∇ y k). By differentiating, applying Minkowski's inequality to the resulting sum of tensors, and applying the Cauchy-Schwarz inequality to each term, we have Using the symmetry of the kernel to combine like terms, this yields that where B < ∞ such that B ≥ |k| + ∇ x k + ∇ x ∇ y k . Since f N (θ) = 0 when N = 1 by definition, we can assume without loss of generality that N ≥ 2, so 1 N −1 = 1 N (1+ 1 N −1 ) ≤ 2/N . Since each term is non-negative, we can add in the i = j terms, By assumption, 1 N i ∇ 2 θ s q θ (X (i) ) : N ∈ N, θ ∈ E is bounded with probability 1, and similarly for 1 N i ∇ 3 θ s q θ (X (i) ) : N ∈ N, θ ∈ E . We show the same for 1 N i s q θ (X (i) ) and 1 N i ∇ θ s q θ (X (i) ) . By Equation 40, we have sup θ∈Ē s q θ (x) P 0 (dx) ≤ g 0,Ē (x)P 0 (dx) < ∞.
Hence, by Theorem 1.3.3 of Ghosh and Ramamoorthi (2003), 1 N i s q θ (X (i) ) converges uniformly onĒ, almost surely. In particular, 1 N i s q θ (X (i) ) is uniformly bounded on E, almost surely. The same argument holds for 1 N i ∇ θ s q θ (X (i) ) using g 1,Ē (x). Therefore, by Equation S44, it follows that θ u θ (X (i) , X (j) ) is uniformly bounded on E. Since k is positive by assumption, i =j k(X (i) , X (j) ) > 0 for all N ≥ 2 and by Equations S39 and S40, 1 N (N −1) i =j k(X (i) , X (j) ) converges a.s. to a finite quantity greater than 0. We conclude that almost surely, is uniformly bounded on E, for N ∈ {2, 3, . . .}. Recall that for N = 1, f N (θ) = 0 by definition. Therefore, almost surely, (f N : N ∈ N) is uniformly bounded on E.
Proof of Theorem 6.9. We show that the conditions of Theorem 3.2 of Miller (2021) are met, from which the conclusions of this theorem follow immediately.
By Condition 6.10 and Equation 35, f N has continuous third-order partial derivatives on Θ. Let E be the set from Condition 6.10. With probability 1, f N → f uniformly on E (by Proposition S3.2 with C =Ē) and (f N ) is uniformly bounded on E (by Proposition S3.3). Note that f is finite on Θ by Proposition 6.3. Thus, by Theorem 3.4 of , f and f exist on E and f N → f uniformly on E with probability 1. Since θ * is a minimizer of f and θ * ∈ E, we know that f (θ * ) = 0 and f (θ * ) is positive semidefinite; thus, f (θ * ) is positive definite since it is invertible by assumption.
Case (a): Now, consider the case where Θ is compact. Then almost surely, f N → f uniformly on Θ by Proposition S3.2 with C = Θ. Since θ * is a unique minimizer of f , we By uniform convergence, with probability 1, there exists N such that for all N > N , sup θ∈Θ |f N (θ) − f (θ)| ≤ /2, and thus, Hence, lim inf N inf θ∈Θ\H f N (θ) > f (θ * ) almost surely. Applying Theorem 3.2 of , the conclusion of the theorem follows. Note that f N (θ N ) → f (θ * ) a.s. since θ N → θ * and f N → f uniformly on E.
Case (b): Alternatively, consider the case where Θ is open and f N is convex on Θ. For all θ ∈ Θ, with probability 1, f N (θ) → f (θ) (by Proposition S3.2 with C = {θ}). However, we need to show that with probability 1, for all θ ∈ Θ, f N (θ) → f (θ). We follow the argument in the proof of Theorem 6.3 of . Let W be a countable dense subset of Θ. Since W is countable, with probability 1, for all θ ∈ W , f N (θ) → f (θ). Since f N is convex, then with probability 1, for all θ ∈ Θ, the limitf (θ) := lim N f N (θ) exists and is finite, andf is convex (Theorem 10.8 of Rockafellar, 1970). Since f N is convex and f (θ) is finite, f (θ) is also convex. Since f andf are convex, they are also continuous (Theorem 10.1 of Rockafellar, 1970). Continuous functions that agree on a dense subset of points must be equal. Thus, with probability 1, for all θ ∈ Θ, f N (θ) → f (θ). Applying Theorem 3.2 of , the conclusion of the theorem follows.
The denominator converges a.s. to a finite positive constant, as in the proof of Proposition S3.2. It is straightforward to verify that E X,Y ∼P 0 [ ∇ θ u θ * (X, Y ) 2 ] < ∞ since s q θ * and ∇ θ θ=θ * s q θ are in L 2 (P 0 ) by assumption. By Theorems 5.5.1A and 5.5.2 of Serfling (2009) Further, by the Leibniz integral rule (Folland, 1999, Theorem 2.27), using the fact that f (θ * ) = 0 since θ * is a minimizer of f . Thus, Next, we examine the convergence of θ N to θ * . For all N sufficiently large, f N (θ N ) = 0 by Theorem 6.9 (part 1), and thus, by Taylor's theorem, where θ + N is on the line between θ N and θ * . As in the proof of Theorem 6.9, f N → f uniformly on the set E defined in Condition 6.10. Thus, since f N is continuous on E and In particular, f N (θ + N ) is invertible for all N sufficiently large, since f (θ * ) is invertible by assumption. Hence, and therefore, by Equation S47, This result matches Theorem 4 in . By Taylor's theorem, for all N sufficiently large, where θ ++ N is on the line between θ N and θ * . Therefore, using the same reasoning as for Equations S48 and S50,
It is straightforward to verify that E X,Y ∼P 0 [|u θ * (X, Y )| 2 ] < ∞ since s q θ * is in L 2 (P 0 ). By Theorems 5.5.1A and 5.5.2 of Serfling (2009) Similarly, since k is bounded, It is straightforward to check that the second part of the theorem (Equation 44) follows.
where C j is a constant that does not depend on N . Hence, log K 1,N K 2,N + N (f 1,N (θ 1,N ) − f 2,N (θ 2,N )) By Theorem 6.12, f j,N (θ j,N ) P 0 − → f j (θ j, * ), and therefore, Plugging in the definition of f j (Equation 36), this proves part 1 of the theorem. For part 2, suppose f 1 (θ 1, * ) = f 2 (θ 2, * ) = 0 and m B 2 − m B 1 does not depend on N . Then by Theorem 6.12, f j,N (θ j,N ) = O P 0 (N −1 ). Using this in Equation S54, we have For part 3, suppose f 1 (θ 1, * ) = f 2 (θ 2, * ) and m B j = c B j √ N . Then by Theorem 6.12, f j,N (θ j,N ) = f j (θ j, * ) + O P 0 (N −1/2 ). Using this in Equation S54, we have S4 Additional probabilistic PCA details S4.1 Optimizing the NKSD Computing the Laplace or BIC approximation to the SVC requires finding the minimizer of nksd(p 0 (x) q(x|θ)) with respect to θ. In this section, we describe how components of the nksd can be pre-computed to speed up this optimization process. The generative model used for pPCA can be rewritten using the properties of multivariate normal distributions as X ∼ N (0, HH + vI d ).
The Stein score function for the pPCA model is then s q θ (x) = ∇ x log q(x|H, v) = −(HH + vI d ) −1 x.
Define the matrices K ij := I(i = j) k(X (i) , X (j) ), where I(E) is the indicator function, which equals 1 when E is true and is 0 otherwise. Define the scalarsK := (X (i) , X (j) ).
Letting X ∈ R N ×d be the data matrix, the NKSD can be written as nksd(p 0 (x) q(x|H, v)) = 1 K trace(X KX(HH + vI d ) −1 (HH + vI d ) −1 ) − 2 trace(X K (HH + vI d ) −1 ) +K , where we have used the fact that the kernel is symmetric. The terms X KX and X K are the only ones that include sums over the entire dataset; these can be pre-computed, before optimizing the parameters H and v.
To compute the matrix inversion (HH + vI d ) −1 we follow the strategy of Minka (2001) Thus, applying the Woodbury matrix identity and using I d U = U = U I k I k = U I k U U , Therefore, (HH + vI d ) −1 = U (L −1 − v −1 I k )U + v −1 I d .

S23
Computing L −1 is trivial since the matrix is diagonal. Returning to the nksd we have nksd(p 0 (x) q(x|U, L, v)) We optimized U , L and v using the trust region method implemented in pymanopt .

S4.2 Data selection with the SVC
We used the approximate optimum technique in Section 2.3.3 to estimate the SVC for different foreground subspaces. Following Section S1.2, we used the factored IMQ kernel with β = −0.5 and c = 1. We focused on foreground subspaces that correspond to subsets of the data dimensions. More specifically, recall that X F = V X; then, we impose the restriction that each column of V is a standard basis vector e (b) ∈ R d , where e where H S F is the submatrix consisting of rows of H indexed by S F , and |S F | is the size of the set S F . In the projected model, some of the parameters are nuisance variables with no contribution to the likelihood. Since the dimension of a d × k matrix on the Stiefel manifold is dk − k(k + 1)/2, the total dimension of the foreground model (including contributions from parameters U , L and v) is m F = |S F |k − k(k + 1)/2 + k + 1, assuming |S F | ≥ k.

S4.3 Calibration
The T hyperparameter was calibrated as in Section S1.1. In detail, we sampled 10 independent true parameter values from the prior, with α = 1 and d = 6. (We used a slightly less disperse prior than during inference, where we set α = 0.1, to avoid numerical instabilities in theT estimate.) Then, for each of the true parameter values, we simulated N = 2000 datapoints. For each simulated true parameter value, we tracked the trend in theT estimator S24 T < l a t e x i t s h a 1 _ b a s e 6 4 = " J E U Q S w X j R 0 M 7 j i u k 8 V M X 9 7 W d A H 4 = " > A A A B 7 n i c b V B N S 8 N A E J 3 U r 1 q / q h 6 9 L B b B U 0 m q o M e i F 4 8 V + g V t K J v t p l 2 6 2 Y T d i V B C f 4 Q X D 4 p 4 9 f d 4 8 9 + 4 b X P Q 1 g c D j / d m m J k X J F I Y d N 1 v p 7 C x u b W 9 U 9 w t 7 e 0 f H B 6 V j 0 / a J k 4 1 4 y 0 W y 1 h 3 A 2 q 4 F I q 3 U K D k 3 U R z G g W S d 4 L J / d z v P H F t R K y a O E 2 4 H 9 G R E q F g F K 3 U 6 Y 8 p Z s 3 Z o F x x q + 4 C Z J 1 4 O a l A j s a g / N U f x i y N u E I m q T E 9 z 0 3 Q z 6 h G w S S f l f q p 4 Q l l E z r i P U s V j b j x s 8 W 5 M 3 J h l S E J Y 2 1 L I V m o v y c y G h k z j Q L b G V E c m 1 V v L v 7 n 9 V I M b / 1 M q C R F r t h y U Z h K g j G Z / 0 6 G Q n O G c m o J Z V r Y W w k b U 0 0 Z 2 o R K N g R v 9 e V 1 0 q 5 V v a t q 7 f G 6 U r / L 4 y j C G Z z D J X h w A 3 V 4 g A a 0 g M E E n u E V 3 p z E e X H e n Y 9 l a 8 H J Z 0 7 h D 5 z P H 3 r W j 6 k = < / l a t e x i t > Figure S2: Estimated T for increasing number of data samples, for 10 independent parameter samples from the prior. The median value at N = 2000 isT = 0.052.
(Equation S1) with increasing N ( Figure S2). The median estimated T value at N = 2000 was 0.052 across the 10 runs.

S4.5 Datasets and preprocessing
We downloaded two publicly available datasets. The first dataset was taken from human peripheral blood mononuclear cells (PBMCs): https://support.10xgenomics.com/ single-cell-gene-expression/datasets/1.1.0/pbmc3k. This is a standard dataset used in the tutorials for Seurat (Stuart et al., 2019) and Scanpy , for example. The second was taken from a dissociated extranodal marginal zone B-cell tumor, specifically a mucosa-associated lymphoid tissue (MALT) tumor: https://support.10xgenomics.com/ single-cell-gene-expression/datasets/3.0.0/malt_10k_protein_v3. We pre-processed the data using Scprep (Gigante et al., 2020), following its example: we normalized the total expression of each cell to match the median total expression in the dataset, to account for variability in library size, and then square-root transformed the resulting normalized counts.
For posterior inference, we employ a mean-field variational approximation: independent normal distributions for the entries of H j , normal distributions forμ andτ , and Laplace  Figure S5: Comparison of the 187 selected genes and 13 excluded genes using data selection.
(a) Violin plot ofσ j over all excluded and selected genes j, respectively, when applying the model to all 200 genes, whereσ j is the mean posterior standard deviation of the interaction energies ∆E jj for gene j, that is,σ j := 1 d−1 j =j std(∆E jj | data). (b) Violin plot of f j over all excluded and selected genes j, respectively, where f j is the fraction of cells with count equal to zero for gene j. The data selection procedure excluded all genes with more than 85% zeros and selected all genes with fewer than 85% zeros.